#include <include/frameProcessor.h>

FrameProcessor::FrameProcessor(const string &pb_file_path, const string &pbtxt_file_path)
    : net(cv::dnn::readNetFromTensorflow(pb_file_path, pbtxt_file_path)) {}

void FrameProcessor::processFrame(const cv::Mat &frame, cv::Mat &out_frame)
{
    Mat blob = cv::dnn::blobFromImage(frame, 1.0, cv::Size(300, 300), cv::Scalar(104, 117, 123), false, false);
    net.setInput(blob);
    Mat probs = net.forward(); // 推理
    Mat detectMat(probs.size[2], probs.size[3], CV_32F, probs.ptr<float>());
    out_frame = frame.clone(); // 复制输入帧以绘制结果
    for (int i = 0; i < detectMat.rows; i++)
    {
        float confidence = detectMat.at<float>(i, 2);
        if (confidence > 0.5)
        {
            int x1 = static_cast<int>(detectMat.at<float>(i, 3) * frame.cols);
            int y1 = static_cast<int>(detectMat.at<float>(i, 4) * frame.rows);
            int x2 = static_cast<int>(detectMat.at<float>(i, 5) * frame.cols);
            int y2 = static_cast<int>(detectMat.at<float>(i, 6) * frame.rows);
            Rect rect(x1, y1, x2 - x1, y2 - y1);
            rectangle(out_frame, rect, Scalar(0, 255, 0), 2, 8);
        }
    }
}